import torch
import librosa
import numpy as np

def get_numpy_spectrum(x):
    s = librosa.core.stft(x, n_fft=2048, win_length=2048, hop_length=512)
    a = np.abs(s)**2
    feat = librosa.power_to_db(a)
    return feat

# 1025,127
def get_torch_spectrum(x):
    def torch_power_to_db(S, amin=1e-10, top_db=80.0):
        magnitude = S
        log_spec = 10.0 * torch.log10(torch.maximum(torch.Tensor([amin]), magnitude))
        log_spec = torch.maximum(log_spec, log_spec.max() - top_db)
        return log_spec
    x=torch.stft(x,n_fft=2048,win_length=2048,hop_length=512,
                 window=torch.hann_window(2048),return_complex=True)
    x=torch.abs(x)**2
    x=torch_power_to_db(x)
    return x

# 513,253
def get_torch_spectrum_1024(x):
    def torch_power_to_db(S, amin=1e-10, top_db=80.0):
        magnitude = S
        log_spec = 10.0 * torch.log10(torch.maximum(torch.Tensor([amin]), magnitude))
        log_spec = torch.maximum(log_spec, log_spec.max() - top_db)
        return log_spec
    x=torch.stft(x,n_fft=1024,win_length=1024,hop_length=256,
                 window=torch.hann_window(1024),return_complex=True)
    x=torch.abs(x)**2
    x=torch_power_to_db(x)
    return x

#
def get_torch_spectrum_new(x):
    def torch_power_to_db(S, amin=1e-10, top_db=80.0):
        magnitude = S
        log_spec = 10.0 * torch.log10(torch.maximum(torch.Tensor([amin]), magnitude))
        log_spec = torch.maximum(log_spec, log_spec.max() - top_db)
        return log_spec
    x=torch.stft(x,n_fft=2048,win_length=1024,hop_length=256,
                 window=torch.hann_window(1024),return_complex=True)
    x=torch.abs(x)**2
    x=torch_power_to_db(x)
    return x

def get_torch_spectrum_3072(x):
    def torch_power_to_db(S, amin=1e-10, top_db=80.0):
        magnitude = S
        log_spec = 10.0 * torch.log10(torch.maximum(torch.Tensor([amin]), magnitude))
        log_spec = torch.maximum(log_spec, log_spec.max() - top_db)
        return log_spec
    x=torch.stft(x,n_fft=3072,win_length=1024,hop_length=256,
                 window=torch.hann_window(1024),return_complex=True)
    x=torch.abs(x)**2
    x=torch_power_to_db(x)
    return x
